from langchain.agents import create_react_agent, AgentExecutor
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate
from .tools.data_collection_tool import data_collection_tool
from .tools.similarity_calculation_tool import similarity_calculation_tool
from .tools.area_recommendation_tool import area_recommendation_tool
import json
import os

class PeakDispatchAgent:
    def __init__(self):
        # 使用环境变量或默认配置
        self.llm = ChatOpenAI(
            temperature=0.1, 
            model="qwen-max",
            api_key=os.getenv('DASHSCOPE_API_KEY'),
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )
    
        self.tools = [
            data_collection_tool,
            similarity_calculation_tool,
            area_recommendation_tool
        ]

        # 创建ReAct提示模板
        template = """你是一个智能调度助手，专门负责分析司机行为数据并推荐最佳接单区域。

你有以下工具可以使用：
{tools}

使用以下格式：

Question: 需要解决的问题
Thought: 你应该思考要做什么
Action: 要采取的行动，应该是 [{tool_names}] 中的一个
Action Input: 行动的输入
Observation: 行动的结果
... (这个 Thought/Action/Action Input/Observation 可以重复N次)
Thought: 我现在知道最终答案了
Final Answer: 对原始问题的最终答案

工作流程：
1. 首先使用data_collection_tool收集司机历史数据
2. 然后使用similarity_calculation_tool找出相似司机
3. 最后使用area_recommendation_tool推荐最佳区域

开始！

Question: {input}
Thought: {agent_scratchpad}"""

        prompt = PromptTemplate.from_template(template)

        try:
            # 创建ReAct Agent
            agent = create_react_agent(self.llm, self.tools, prompt)
            self.agent = AgentExecutor(agent=agent, tools=self.tools, verbose=True)
        except Exception as e:
            print(f"Agent初始化失败: {e}")
            # 如果Agent初始化失败，使用简单的工具链
            self.agent = None
    
    def execute_dispatch_recommendation(self, driver_id, time_category="morning_peak"):
        """执行完整的调度推荐流程"""
        try:
            if self.agent:
                # 使用LangChain Agent执行
                prompt = f"""
                请为司机ID {driver_id} 在 {time_category} 时段推荐最佳接单区域。
                
                执行步骤：
                1. 收集最近7天的司机接单数据
                2. 计算与目标司机最相似的10个司机
                3. 基于相似司机的经验推荐前5个最佳区域
                
                请提供详细的分析过程和推荐理由。
                """
                
                result = self.agent.invoke({"input": prompt})
                return result.get("output", "未获取到结果")
            else:
                # 直接使用工具链执行
                return self._execute_tools_directly(driver_id, time_category)
                
        except Exception as e:
            return f"调度推荐执行失败：{str(e)}"
    
    def _execute_tools_directly(self, driver_id, time_category):
        """直接使用工具链执行（不依赖LLM）"""
        results = {}

        # 步骤1：数据收集
        data_query = json.dumps({"days": 7, "output_file": "driver_order_data.csv"})
        data_result = data_collection_tool.invoke(data_query)
        results['data_collection'] = data_result

        # 步骤2：相似度计算
        similarity_query = json.dumps({
            "input_file": "driver_order_data.csv",
            "target_driver_id": driver_id,
            "top_k": 10
        })
        similarity_result = similarity_calculation_tool.invoke(similarity_query)
        results['similarity_calculation'] = similarity_result

        # 解析相似司机结果
        try:
            # 从结果中提取相似司机数据
            import re
            similarity_match = re.search(r'相似度计算完成：(.+)', similarity_result, re.DOTALL)
            if similarity_match:
                similarity_data = json.loads(similarity_match.group(1))
                similar_drivers = similarity_data.get('similar_drivers', [])
            else:
                similar_drivers = []
        except:
            similar_drivers = []

        # 步骤3：区域推荐
        recommendation_query = json.dumps({
            "input_file": "driver_order_data.csv",
            "similar_drivers": similar_drivers,
            "time_category": time_category,
            "top_areas": 5
        })
        recommendation_result = area_recommendation_tool.invoke(recommendation_query)
        results['area_recommendation'] = recommendation_result

        # 整合结果
        final_result = {
            "driver_id": driver_id,
            "time_category": time_category,
            "execution_results": results,
            "summary": "调度推荐流程执行完成"
        }

        return json.dumps(final_result, ensure_ascii=False, indent=2)
